package shell

import (
	. "errors"
	"fmt"
	"github.com/mitchellh/go-homedir"
	"golang.org/x/crypto/ssh"
	"io/ioutil"
	"log"
	"time"
)

type SSHLineType int64

const (
	Password  SSHLineType = 1
	PublicKey SSHLineType = 2
)

type Client struct {
	SSHLineType SSHLineType   //ssh连接方式 1. Password 密码2. PublicKey 密钥
	User        string        //用户名
	Password    string        //密码
	Host        string        //地址
	SSHKeyPath  string        //ssh客户端私钥地址
	LastResult  string        //客户端最后执行结果
	Port        int           //端口
	Timeout     time.Duration //超时
	client      *ssh.Client
}

//创建密码连接客户端
func NewShellPwd(user, pwd, host string, time time.Duration, port int) *Client {
	cli := &Client{
		SSHLineType: Password,
		User:        user,
		Password:    pwd,
		Host:        host,
		Port:        port,
		Timeout:     time,
		client:      nil,
	}
	return cli
}

//创建密钥连接客户端
func NewShellPublicKey(user, keyPath, host string, time time.Duration, port int) *Client {
	cli := &Client{
		SSHLineType: PublicKey,
		User:        user,
		Host:        host,
		SSHKeyPath:  keyPath,
		Port:        port,
		Timeout:     time,
		client:      nil,
	}
	return cli
}

func (cli *Client) Run(cmd string) (string, error) {
	if cli.client == nil {
		if err := cli.connect(); err != nil {
			return "", err
		}
	}
	session, err := cli.client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()
	buf, err := session.CombinedOutput(cmd)

	cli.LastResult = string(buf)
	return cli.LastResult, err
}

//基础连接ssh
func (cli *Client) connect() error {
	config := &ssh.ClientConfig{
		Timeout:         time.Second, //ssh 连接time out 时间一秒钟, 如果ssh验证错误 会在一秒内返回
		User:            cli.User,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //这个可以， 但是不够安全
	}

	switch cli.SSHLineType {
	case Password:
		config.Auth = []ssh.AuthMethod{ssh.Password(cli.Password)}
	case PublicKey:
		config.Auth = []ssh.AuthMethod{publicKeyAuthFunc(cli.SSHKeyPath)}
	default:
		return New("SSHLineType can not be empty!")
	}
	addr := fmt.Sprintf("%s:%d", cli.Host, cli.Port)
	sshClient, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		log.Fatal("创建ssh client 失败", err)
		return err
	}
	cli.client = sshClient
	return nil
}

func (cli *Client) Close() error {
	return cli.client.Close()
}

func publicKeyAuthFunc(kPath string) ssh.AuthMethod {
	keyPath, err := homedir.Expand(kPath)
	if err != nil {
		log.Fatal("find key's home dir failed", err)
	}
	key, err := ioutil.ReadFile(keyPath)
	if err != nil {
		log.Fatal("ssh key file read failed", err)
	}
	// Create the Signer for this private key.
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		log.Fatal("ssh key signer failed", err)
	}
	return ssh.PublicKeys(signer)
}
